# ... existing code ...

import os
import shutil
import zipfile
import tempfile
from openpyxl import load_workbook

def test_split_excel_by_worksheet(input_file, output_dir=None):
    # 设置输出目录
    if output_dir is None:
        input_dir = os.path.dirname(input_file)
        input_filename = os.path.splitext(os.path.basename(input_file))[0]
        output_dir = os.path.join(input_dir, f"{input_filename}_拆分")
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    print(f"正在处理文件: {input_file}")
    print(f"输出目录: {output_dir}")
    
    # 加载工作簿获取工作表名称
    wb = load_workbook(input_file, read_only=True)
    sheet_names = wb.sheetnames
    wb.close()
    
    # 创建临时目录
    with tempfile.TemporaryDirectory() as temp_dir:
        # 解压原始Excel文件
        with zipfile.ZipFile(input_file, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # 检查是否有媒体文件
        media_dir = os.path.join(temp_dir, "xl", "media")
        has_media = os.path.exists(media_dir)
        if has_media:
            print(f"发现媒体文件夹，包含 {len(os.listdir(media_dir))} 个文件")
        
        # 处理每个工作表
        for sheet_index, sheet_name in enumerate(sheet_names, 1):
            print(f"正在处理工作表 {sheet_index}/{len(sheet_names)}: {sheet_name}")
            
            # 为每个工作表创建一个新的Excel文件
            # 首先创建一个只有一个工作表的工作簿
            temp_wb = load_workbook(input_file)
            
            # 保留当前工作表，删除其他工作表
            target_sheet = temp_wb[sheet_name]
            for s in list(temp_wb.sheetnames):
                if s != sheet_name:
                    temp_wb.remove(temp_wb[s])
            
            # 保存临时文件
            output_file = os.path.join(output_dir, f"{sheet_name}.xlsx")
            temp_wb.save(output_file)
            temp_wb.close()
            
            # 现在我们有了一个只包含单个工作表的Excel文件
            # 如果原始文件有媒体文件，我们需要确保它们被正确复制
            if has_media:
                # 创建另一个临时目录来处理这个特定的工作表
                sheet_temp_dir = os.path.join(temp_dir, f"sheet_{sheet_index}")
                os.makedirs(sheet_temp_dir, exist_ok=True)
                
                # 解压新创建的Excel文件
                with zipfile.ZipFile(output_file, 'r') as zip_ref:
                    zip_ref.extractall(sheet_temp_dir)
                
                # 检查工作表是否有对应的绘图文件
                drawing_file = os.path.join(temp_dir, "xl", "drawings", f"drawing{sheet_index}.xml")
                drawing_rels_file = os.path.join(temp_dir, "xl", "drawings", "_rels", f"drawing{sheet_index}.xml.rels")
                
                if os.path.exists(drawing_file) and os.path.exists(drawing_rels_file):
                    print(f"工作表 {sheet_name} 有绘图文件")
                    
                    # 确保目标目录存在
                    target_drawings_dir = os.path.join(sheet_temp_dir, "xl", "drawings")
                    target_drawings_rels_dir = os.path.join(target_drawings_dir, "_rels")
                    os.makedirs(target_drawings_dir, exist_ok=True)
                    os.makedirs(target_drawings_rels_dir, exist_ok=True)
                    
                    # 复制绘图文件
                    shutil.copy(drawing_file, os.path.join(target_drawings_dir, "drawing1.xml"))
                    shutil.copy(drawing_rels_file, os.path.join(target_drawings_rels_dir, "drawing1.xml.rels"))
                    
                    # 复制媒体文件
                    target_media_dir = os.path.join(sheet_temp_dir, "xl", "media")
                    os.makedirs(target_media_dir, exist_ok=True)
                    
                    # 从关系文件中读取媒体文件引用
                    with open(drawing_rels_file, 'r', encoding='utf-8') as f:
                        rels_content = f.read()
                    
                    # 复制所有引用的媒体文件
                    import re
                    media_files = re.findall(r'Target="../media/([^"]+)"', rels_content)
                    for media_file in media_files:
                        src_media_path = os.path.join(media_dir, media_file)
                        if os.path.exists(src_media_path):
                            shutil.copy(src_media_path, os.path.join(target_media_dir, media_file))
                            print(f"已复制媒体文件: {media_file}")
                    
                    # 更新工作表与绘图的关系
                    sheet_rels_dir = os.path.join(sheet_temp_dir, "xl", "worksheets", "_rels")
                    os.makedirs(sheet_rels_dir, exist_ok=True)
                    
                    # 复制工作表关系文件
                    sheet_rels_file = os.path.join(temp_dir, "xl", "worksheets", "_rels", f"sheet{sheet_index}.xml.rels")
                    if os.path.exists(sheet_rels_file):
                        shutil.copy(sheet_rels_file, os.path.join(sheet_rels_dir, "sheet1.xml.rels"))
                        
                        # 修改关系文件中的引用，将drawing{sheet_index}.xml改为drawing1.xml
                        with open(os.path.join(sheet_rels_dir, "sheet1.xml.rels"), 'r', encoding='utf-8') as f:
                            content = f.read()
                        
                        content = content.replace(f'drawing{sheet_index}.xml', 'drawing1.xml')
                        
                        with open(os.path.join(sheet_rels_dir, "sheet1.xml.rels"), 'w', encoding='utf-8') as f:
                            f.write(content)
                    
                    # 更新 [Content_Types].xml 确保包含绘图和媒体类型
                    content_types_file = os.path.join(sheet_temp_dir, "[Content_Types].xml")
                    with open(content_types_file, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    # 添加必要的内容类型
                    if '<Default Extension="png" ContentType="image/png"/>' not in content:
                        content = content.replace('</Types>', '<Default Extension="png" ContentType="image/png"/></Types>')
                    
                    if '<Default Extension="jpeg" ContentType="image/jpeg"/>' not in content:
                        content = content.replace('</Types>', '<Default Extension="jpeg" ContentType="image/jpeg"/></Types>')
                    
                    if '<Default Extension="jpg" ContentType="image/jpeg"/>' not in content:
                        content = content.replace('</Types>', '<Default Extension="jpg" ContentType="image/jpeg"/></Types>')
                    
                    if '<Override PartName="/xl/drawings/drawing1.xml" ContentType="application/vnd.openxmlformats-officedocument.drawing+xml"/>' not in content:
                        content = content.replace('</Types>', '<Override PartName="/xl/drawings/drawing1.xml" ContentType="application/vnd.openxmlformats-officedocument.drawing+xml"/></Types>')
                    
                    with open(content_types_file, 'w', encoding='utf-8') as f:
                        f.write(content)
                    
                    # 重新打包Excel文件
                    os.remove(output_file)  # 删除之前的文件
                    with zipfile.ZipFile(output_file, 'w') as zip_out:
                        for root, dirs, files in os.walk(sheet_temp_dir):
                            for file in files:
                                file_path = os.path.join(root, file)
                                arc_name = os.path.relpath(file_path, sheet_temp_dir)
                                zip_out.write(file_path, arc_name)
                    
                    print(f"已重新打包工作表 {sheet_name} 的Excel文件，包含图片")
            
            print(f"工作表 {sheet_name} 已保存到 {output_file}")
    
    print("所有工作表拆分完成")

if __name__ == "__main__":
    # 使用相对路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    input_file = os.path.join(current_dir, "拆分例子.xlsx")
    
    # 输出目录将自动创建为"月报绘图_拆分"
    test_split_excel_by_worksheet(input_file)